#pragma once

#include <functional>
#include <unordered_map>

class RBTreeRankBase
{
public:
    RBTreeRankBase();
    ~RBTreeRankBase();

protected:
    enum Color{ RED, BLACK };
    struct NodeBase {
        Color color;
        NodeBase *p;
        NodeBase *left;
        NodeBase *right;
        size_t n;
    };

    void LeftRotate(NodeBase *x);
    void RightRotate(NodeBase *x);

    void RBTransplant(NodeBase *u, NodeBase *v);

    void RBInsertFixup(NodeBase *z);

    void RBDelete(NodeBase *z);
    void RBDeleteFixup(NodeBase *x);

    NodeBase *TreeMinimum(NodeBase *x) const;
    NodeBase *TreeMaximum(NodeBase *x) const;

    NodeBase *TreeSuccessor(NodeBase *x) const;
    NodeBase *TreePredecessor(NodeBase *x) const;

    size_t TreeRank(NodeBase *x) const;
    NodeBase *TreeAdvance(size_t i) const;
    void TreeNFixup(NodeBase *x, ssize_t n) const;

    NodeBase *root_;
    NodeBase nil_;
};

template <class Key, class T, class Compare = std::less<T>>
class RBTreeRank : public RBTreeRankBase
{
public:
    struct Node : public NodeBase {
        Key key;
        T v;
    };
    class Iterator {
    public:
        Iterator(const RBTreeRank *rank, Node *node);
        Iterator &operator++();
        Iterator &operator--();
        const Node &operator*() const;
        const Node *operator->() const;
        bool operator==(Iterator other) const;
        bool operator!=(Iterator other) const;
    private:
        const RBTreeRank *rank;
        Node *node;
    };

    RBTreeRank() {}
    RBTreeRank(Compare &&compare)
        : rank_compare_(std::move(compare)) {}
    ~RBTreeRank() { Clear(); }

    bool Add(Key key, T &&v);
    bool Remove(Key key);
    void Clear();

    template <typename F>
    bool Dirty(Key key, F &&f);

    size_t GetRank(Key key) const;

    Iterator find(Key key) const;
    Iterator advance(size_t i) const;
    Iterator begin() const;
    Iterator end() const;

private:
    void RBInsert(NodeBase *z);

    bool Less(const NodeBase *n1, const NodeBase *n2) const;

    const Compare rank_compare_;
    std::unordered_map<Key, Node*> rank_objs_;
};

#include "RBTreeRank.inl"
